"""
Async Content Generation Service
异步内容生成服务 - 使用Python原生asyncio实现异步内容生成
"""

import asyncio
import aiohttp
import uuid
from typing import Optional, Dict, Any, List, Callable
from datetime import datetime, timedelta
import structlog

from ..models.user import User
from ..models.creation_request import CreationRequest, RequestStatus
from ..models.generated_content import GeneratedContent, ContentStatus
from ..models.task_queue import TaskQueue, TaskStatus
from ..core.config import settings
from ..core.exceptions import (
    ContentGenerationError,
    AIModelNotAvailableError,
    RateLimitExceededError,
    ValidationError,
)
from .ai_model_service import AIModelService

# 配置结构化日志
logger = structlog.get_logger(__name__)


class AsyncContentGenerationService:
    """异步内容生成服务类"""
    
    def __init__(self):
        self.ai_service = AIModelService()
        self.active_tasks: Dict[str, asyncio.Task] = {}
        self.progress_callbacks: Dict[str, List[Callable]] = {}
        self.completion_callbacks: Dict[str, List[Callable]] = {}
        
    async def generate_content_async(
        self,
        creation_request_id: uuid.UUID,
        ai_service: AIModelService,
        content_index: int = 1,
        total_contents: int = 1
    ) -> Optional[GeneratedContent]:
        """
        异步生成内容
        
        Args:
            creation_request_id: 创作请求ID
            ai_service: AI模型服务
            content_index: 内容索引（用于批量生成）
            total_contents: 总内容数量
            
        Returns:
            Optional[GeneratedContent]: 生成的内容对象，失败时返回None
        """
        task_id = str(uuid.uuid4())
        
        try:
            logger.info(
                "Starting async content generation",
                task_id=task_id,
                creation_request_id=str(creation_request_id),
                content_index=content_index,
                total_contents=total_contents
            )
            
            # 获取创作请求详情
            creation_request = await self._get_creation_request(creation_request_id)
            if not creation_request:
                raise ContentGenerationError("创作请求不存在")
            
            # 更新任务状态为处理中
            await self._update_task_status(creation_request_id, TaskStatus.PROCESSING, 10)
            
            # 验证用户权限和限制
            user = creation_request.user
            if not user.can_generate_content():
                raise RateLimitExceededError("用户生成限制已达到或存储空间已满")
            
            # 进度回调
            await self._notify_progress(creation_request_id, 20, "开始生成内容...")
            
            # 准备生成参数
            generation_params = self._prepare_generation_params(creation_request, content_index)
            
            # 异步调用AI模型
            await self._notify_progress(creation_request_id, 30, "调用AI模型...")
            
            start_time = datetime.utcnow()
            
            # 使用asyncio.create_task实现真正的异步调用
            ai_task = asyncio.create_task(
                self._call_ai_model_async(
                    ai_service,
                    creation_request.prompt,
                    creation_request.content_type,
                    creation_request.ai_provider,
                    generation_params
                )
            )
            
            # 等待AI模型响应，设置超时
            try:
                ai_result = await asyncio.wait_for(ai_task, timeout=300)  # 5分钟超时
            except asyncio.TimeoutError:
                raise ContentGenerationError("AI模型响应超时")
            
            processing_time = (datetime.utcnow() - start_time).total_seconds()
            
            await self._notify_progress(creation_request_id, 70, "AI模型调用完成，开始处理结果...")
            
            # 处理AI生成结果
            generated_content = await self._process_ai_result(
                ai_result,
                creation_request,
                content_index,
                processing_time
            )
            
            await self._notify_progress(creation_request_id, 90, "内容生成完成...")
            
            # 更新任务状态为完成
            await self._update_task_status(creation_request_id, TaskStatus.COMPLETED, 100)
            
            # 通知完成回调
            await self._notify_completion(creation_request_id, generated_content)
            
            logger.info(
                "Async content generation completed successfully",
                task_id=task_id,
                creation_request_id=str(creation_request_id),
                content_id=str(generated_content.id) if generated_content else None,
                processing_time=processing_time
            )
            
            return generated_content
            
        except Exception as e:
            logger.error(
                "Async content generation failed",
                task_id=task_id,
                creation_request_id=str(creation_request_id),
                error=str(e),
                error_type=type(e).__name__
            )
            
            # 更新任务状态为失败
            await self._update_task_status(
                creation_request_id, 
                TaskStatus.FAILED, 
                0,
                error_message=str(e)
            )
            
            # 通知错误回调
            await self._notify_error(creation_request_id, e)
            
            return None
            
        finally:
            # 清理任务记录
            self._cleanup_task(task_id)
    
    async def generate_content_batch_async(
        self,
        creation_request_id: uuid.UUID,
        ai_service: AIModelService,
        count: int = 1
    ) -> List[Optional[GeneratedContent]]:
        """
        批量异步生成内容
        
        Args:
            creation_request_id: 创作请求ID
            ai_service: AI模型服务
            count: 生成数量
            
        Returns:
            List[Optional[GeneratedContent]]: 生成的内容列表
        """
        logger.info(
            "Starting batch async content generation",
            creation_request_id=str(creation_request_id),
            count=count
        )
        
        # 使用asyncio.gather并发执行多个生成任务
        tasks = []
        for i in range(count):
            task = self.generate_content_async(
                creation_request_id,
                ai_service,
                content_index=i + 1,
                total_contents=count
            )
            tasks.append(task)
        
        # 并发执行所有任务
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        successful_results = []
        failed_count = 0
        
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error(
                    f"Batch generation task {i+1} failed",
                    creation_request_id=str(creation_request_id),
                    error=str(result)
                )
                failed_count += 1
            elif result is not None:
                successful_results.append(result)
            else:
                failed_count += 1
        
        logger.info(
            "Batch async content generation completed",
            creation_request_id=str(creation_request_id),
            successful_count=len(successful_results),
            failed_count=failed_count,
            total_count=count
        )
        
        return successful_results
    
    async def generate_content_with_retry_async(
        self,
        creation_request_id: uuid.UUID,
        ai_service: AIModelService,
        max_retries: int = 3,
        retry_delay: float = 2.0
    ) -> Optional[GeneratedContent]:
        """
        带重试机制的异步内容生成
        
        Args:
            creation_request_id: 创作请求ID
            ai_service: AI模型服务
            max_retries: 最大重试次数
            retry_delay: 重试延迟（秒）
            
        Returns:
            Optional[GeneratedContent]: 生成的内容对象
        """
        for attempt in range(max_retries + 1):
            try:
                logger.info(
                    f"Attempting content generation (attempt {attempt + 1}/{max_retries + 1})",
                    creation_request_id=str(creation_request_id)
                )
                
                result = await self.generate_content_async(
                    creation_request_id,
                    ai_service
                )
                
                if result is not None:
                    return result
                
                # 如果没有返回内容，可能是临时失败，继续重试
                if attempt < max_retries:
                    logger.warning(
                        f"Content generation failed, retrying in {retry_delay} seconds",
                        creation_request_id=str(creation_request_id),
                        attempt=attempt + 1
                    )
                    await asyncio.sleep(retry_delay)
                
            except Exception as e:
                logger.error(
                    f"Content generation attempt {attempt + 1} failed",
                    creation_request_id=str(creation_request_id),
                    error=str(e)
                )
                
                if attempt < max_retries:
                    await asyncio.sleep(retry_delay)
        
        logger.error(
            "All content generation attempts failed",
            creation_request_id=str(creation_request_id),
            max_retries=max_retries
        )
        
        return None
    
    def add_progress_callback(
        self, 
        creation_request_id: uuid.UUID, 
        callback: Callable
    ) -> None:
        """添加进度回调函数"""
        request_id_str = str(creation_request_id)
        if request_id_str not in self.progress_callbacks:
            self.progress_callbacks[request_id_str] = []
        self.progress_callbacks[request_id_str].append(callback)
    
    def add_completion_callback(
        self, 
        creation_request_id: uuid.UUID, 
        callback: Callable
    ) -> None:
        """添加完成回调函数"""
        request_id_str = str(creation_request_id)
        if request_id_str not in self.completion_callbacks:
            self.completion_callbacks[request_id_str] = []
        self.completion_callbacks[request_id_str].append(callback)
    
    async def cancel_generation(self, creation_request_id: uuid.UUID) -> bool:
        """
        取消内容生成任务
        
        Args:
            creation_request_id: 创作请求ID
            
        Returns:
            bool: 是否成功取消
        """
        request_id_str = str(creation_request_id)
        
        # 查找并取消相关的异步任务
        tasks_to_cancel = []
        for task_id, task in self.active_tasks.items():
            if request_id_str in str(task):
                tasks_to_cancel.append(task)
        
        for task in tasks_to_cancel:
            if not task.done():
                task.cancel()
        
        # 等待任务取消完成
        if tasks_to_cancel:
            await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
        
        # 更新任务状态
        await self._update_task_status(
            creation_request_id, 
            TaskStatus.CANCELLED, 
            0
        )
        
        logger.info(
            "Content generation cancelled",
            creation_request_id=str(creation_request_id),
            cancelled_tasks=len(tasks_to_cancel)
        )
        
        return len(tasks_to_cancel) > 0
    
    # 私有辅助方法
    
    async def _get_creation_request(self, request_id: uuid.UUID) -> Optional[CreationRequest]:
        """获取创作请求详情"""
        from sqlalchemy import select
        from ..core.database import get_db
        
        async for db in get_db():
            result = await db.execute(
                select(CreationRequest).where(CreationRequest.id == request_id)
            )
            return result.scalar_one_or_none()
    
    async def _update_task_status(
        self, 
        request_id: uuid.UUID, 
        status: TaskStatus, 
        progress: int,
        error_message: Optional[str] = None
    ) -> None:
        """更新任务状态"""
        from sqlalchemy import select, update
        from ..core.database import get_db
        
        async for db in get_db():
            # 查找相关的任务队列记录
            result = await db.execute(
                select(TaskQueue).where(
                    TaskQueue.related_request_id == request_id
                )
            )
            task_queue = result.scalar_one_or_none()
            
            if task_queue:
                task_queue.status = status
                task_queue.progress = progress
                if error_message:
                    task_queue.error_message = error_message
                
                await db.commit()
    
    def _prepare_generation_params(
        self, 
        creation_request: CreationRequest, 
        content_index: int
    ) -> Dict[str, Any]:
        """准备生成参数"""
        params = creation_request.parameters or {}
        
        # 添加内容索引信息
        params["content_index"] = content_index
        params["total_count"] = creation_request.count
        
        # 设置随机种子（如果有）
        if "seed" not in params and creation_request.seed:
            params["seed"] = creation_request.seed + content_index
        
        return params
    
    async def _call_ai_model_async(
        self,
        ai_service: AIModelService,
        prompt: str,
        content_type: str,
        provider: str,
        parameters: Dict[str, Any]
    ) -> Dict[str, Any]:
        """异步调用AI模型"""
        # 使用aiohttp进行异步HTTP调用
        async with aiohttp.ClientSession() as session:
            # 这里应该调用实际的AI模型API
            # 为了演示，我们模拟异步调用
            await asyncio.sleep(1)  # 模拟网络延迟
            
            # 调用AI服务
            result = await ai_service.generate_content_async(
                prompt=prompt,
                content_type=content_type,
                provider=provider,
                parameters=parameters,
                session=session
            )
            
            return result
    
    async def _process_ai_result(
        self,
        ai_result: Dict[str, Any],
        creation_request: CreationRequest,
        content_index: int,
        processing_time: float
    ) -> GeneratedContent:
        """处理AI生成结果"""
        from sqlalchemy import select
        from ..core.database import get_db
        
        async for db in get_db():
            # 创建生成的内容记录
            generated_content = GeneratedContent(
                user_id=creation_request.user_id,
                creation_request_id=creation_request.id,
                title=ai_result.get("title", f"Generated Content {content_index}"),
                description=creation_request.prompt[:200] if creation_request.prompt else None,
                content_type=creation_request.content_type,
                format=ai_result.get("format", "png"),
                file_name=ai_result.get("file_name", f"generated_{content_index}.png"),
                file_size_bytes=ai_result.get("file_size", 0),
                file_path=ai_result.get("file_path", ""),
                file_url=ai_result.get("file_url", ""),
                width=ai_result.get("width"),
                height=ai_result.get("height"),
                prompt=creation_request.prompt,
                negative_prompt=creation_request.negative_prompt,
                ai_provider=creation_request.ai_provider,
                ai_model=ai_result.get("model"),
                model_version=ai_result.get("model_version"),
                seed=ai_result.get("seed"),
                generation_parameters=creation_request.parameters,
                status=ContentStatus.COMPLETED,
                processing_time_seconds=int(processing_time),
                generated_at=datetime.utcnow()
            )
            
            db.add(generated_content)
            await db.commit()
            await db.refresh(generated_content)
            
            return generated_content
    
    async def _notify_progress(
        self, 
        creation_request_id: uuid.UUID, 
        progress: int, 
        message: str
    ) -> None:
        """通知进度更新"""
        request_id_str = str(creation_request_id)
        callbacks = self.progress_callbacks.get(request_id_str, [])
        
        for callback in callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(progress, message)
                else:
                    callback(progress, message)
            except Exception as e:
                logger.error(
                    "Progress callback failed",
                    request_id=request_id_str,
                    callback=str(callback),
                    error=str(e)
                )
    
    async def _notify_completion(
        self, 
        creation_request_id: uuid.UUID, 
        content: GeneratedContent
    ) -> None:
        """通知完成"""
        request_id_str = str(creation_request_id)
        callbacks = self.completion_callbacks.get(request_id_str, [])
        
        for callback in callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(content)
                else:
                    callback(content)
            except Exception as e:
                logger.error(
                    "Completion callback failed",
                    request_id=request_id_str,
                    callback=str(callback),
                    error=str(e)
                )
    
    async def _notify_error(
        self, 
        creation_request_id: uuid.UUID, 
        error: Exception
    ) -> None:
        """通知错误"""
        request_id_str = str(creation_request_id)
        callbacks = self.completion_callbacks.get(request_id_str, [])
        
        for callback in callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(None, error)
                else:
                    callback(None, error)
            except Exception as e:
                logger.error(
                    "Error callback failed",
                    request_id=request_id_str,
                    callback=str(callback),
                    error=str(e)
                )
    
    def _cleanup_task(self, task_id: str) -> None:
        """清理任务记录"""
        if task_id in self.active_tasks:
            del self.active_tasks[task_id]
    
    async def get_generation_status(self, creation_request_id: uuid.UUID) -> Dict[str, Any]:
        """
        获取生成状态
        
        Args:
            creation_request_id: 创作请求ID
            
        Returns:
            Dict[str, Any]: 状态信息
        """
        request_id_str = str(creation_request_id)
        
        # 检查是否有活跃任务
        has_active_task = any(
            request_id_str in str(task) 
            for task in self.active_tasks.values() 
            if not task.done()
        )
        
        return {
            "has_active_task": has_active_task,
            "active_task_count": len([
                task for task in self.active_tasks.values() 
                if not task.done() and request_id_str in str(task)
            ]),
            "callback_count": len(self.progress_callbacks.get(request_id_str, []))
        }


# 全局异步内容生成服务实例
async_content_service = AsyncContentGenerationService()


# 导出服务
__all__ = ["AsyncContentGenerationService", "async_content_service"]